import ModelInterfaces as mi
import torch
import numpy as np


class NeuralASR(mi.IASRModel):
    word_locations_in_samples = None
    audio_transcript = None

    def __init__(self, model: torch.nn.Module, decoder) -> None:
        super().__init__()
        self.model = model
        self.decoder = decoder  # Decoder from CTC-outputs to transcripts

    def getTranscript(self) -> str:
        """Get the transcripts of the process audio"""
        assert self.audio_transcript is not None, \
               'Can get audio transcripts without having processed the audio'
        return self.audio_transcript

    def getWordLocations(self) -> list:
        """Get the pair of words location from audio"""
        assert self.word_locations_in_samples is not None, \
               'Can get word locations without having processed the audio'
        return self.word_locations_in_samples

    def processAudio(self, audio: torch.Tensor):
        """Process the audio"""
        audio_length_in_samples = audio.shape[1]
        with torch.inference_mode():
            nn_output = self.model(audio)
            # Assuming decoder returns transcript and word_locations
            # This part is specific to Silero's STT output format
            decoded_output = self.decoder(nn_output[0, :, :].detach(), word_align=True) 
            
            # Silero decoder might return a list of dicts or a more complex structure
            # Based on common Silero usage for word timestamps:
            if isinstance(decoded_output, tuple) and len(decoded_output) == 2: # transcript, word_timestamps_list
                self.audio_transcript = decoded_output[0]
                self.word_locations_in_samples = decoded_output[1] # Expects list of dicts with 'start_ts', 'end_ts', 'word'
            elif isinstance(decoded_output, list): # If it directly returns list of word dicts
                self.audio_transcript = " ".join([d['word'] for d in decoded_output if 'word' in d])
                self.word_locations_in_samples = decoded_output
            else: # Fallback if format is unexpected
                print(f"Warning: Silero decoder output format not as expected: {type(decoded_output)}")
                self.audio_transcript = str(decoded_output) # Or try to infer
                self.word_locations_in_samples = []


class NeuralTTS(mi.ITextToSpeechModel):
    def __init__(self, model: torch.nn.Module, sampling_rate: int) -> None:
        super().__init__()
        self.model = model
        self.sampling_rate = sampling_rate

    def getAudioFromSentence(self, sentence: str) -> np.array:
        with torch.inference_mode():
            # Silero TTS model specific call
            audio_transcript = self.model.apply_tts(texts=[sentence],
                                                    sample_rate=self.sampling_rate,
                                                    # device=self.model.device # Ensure TTS model also handles device if necessary
                                                    )[0]
        return audio_transcript.cpu().numpy() # Ensure output is numpy array on CPU


class NeuralTranslator(mi.ITranslationModel):
    def __init__(self, model: torch.nn.Module, tokenizer) -> None:
        super().__init__()
        self.model = model
        self.tokenizer = tokenizer
        # Ensure model is on the correct device if applicable
        # if hasattr(self.model, 'device'):
        #     self.model.to(self.model.device) 

    def translateSentence(self, sentence: str) -> str:
        """Get the transcripts of the process audio"""
        # Ensure model and tokenizer are on the same device as input tensors
        # device = self.model.device if hasattr(self.model, 'device') else torch.device("cpu")
        tokenized_text = self.tokenizer(sentence, return_tensors='pt') # .to(device)
        translation = self.model.generate(**tokenized_text)
        translated_text = self.tokenizer.batch_decode(
            translation, skip_special_tokens=True)[0]

        return translated_text 




